use std::{
    collections::{HashSet, VecDeque}, hash::Hash, usize
};


pub type Color = char;


#[derive(Debug)]
pub struct Palette<T>(Vec<T>);

impl<T: Copy + Eq + Hash> Palette<T> {
    pub fn new_from(rows: Vec<Vec<T>>) -> Palette<T> {
        Palette(rows.concat())
    }

    fn new_const(c: T, size: &PaletteSize) -> Palette<T> {
        Palette(vec![c; size.w * size.h])
    }

    pub fn at(&self, (h, w): (usize, usize), size: &PaletteSize) -> T {
        self.0[h * size.w + w]
    }

    fn paint_at(&mut self, (h, w): (usize, usize), c: T, size: &PaletteSize) {
        self.0[h * size.w + w] = c
    }

    pub fn paint_patch(&self, locs: &Vec<(usize, usize)>, c: T, size: &PaletteSize) -> Palette<T> {
        let mut vec = self.0.clone();
        for loc in locs {
            vec[loc.0 * size.w + loc.1] = c
        }
        Palette(vec)
    }

    pub fn contains(&self, c: &T) -> bool {
        self.0.contains(c)
    }
}


#[derive(Debug)]
pub struct PaletteSize {
    h: usize, w: usize
}

impl PaletteSize {
    pub fn new_from(h: usize, w: usize) -> PaletteSize {
        PaletteSize { h, w }
    }
}


const DIR: [(i64, i64); 4] = [(1, 0), (0, 1), (-1, 0), (0, -1)];
const BLANK: Color = ' ';
/// Step to the 4 adjacent cells of `(x, y)`.
fn step((x, y): (usize, usize), i: usize, size: &PaletteSize)
    -> Option<(usize, usize)> {
    let (x1, y1) = (
        (i64::try_from(x).unwrap() + DIR[i].0),
        (i64::try_from(y).unwrap() + DIR[i].1)
    );
    if 0 <= x1 && x1 < size.h.try_into().unwrap() &&
        0 <= y1 && y1 < size.w.try_into().unwrap() {
        Some((x1.try_into().unwrap(), y1.try_into().unwrap()))
    } else {
        None
    }
}


/// Returns the number of the patches in `palette`.
pub fn view_simple(palette: &Palette<Color>, size: &PaletteSize) -> u32 {
    let mut checked = Palette::new_const(0u32, size);
    let mut n_patch = 0;

    for i in 0..size.h {
        for j in 0..size.w {
            if checked.at((i, j), size) != 0 || palette.at((i, j), size) == BLANK { continue }
            // else
            n_patch += 1;

            let mut queue: VecDeque<(usize, usize)> = VecDeque::new();

            checked.paint_at((i, j), n_patch, size);
            queue.push_back((i, j));

            let color = palette.at((i, j), size);
            while !queue.is_empty() {
                let (x, y) = queue.pop_front().unwrap();

                for i in 0..4 {
                    if let Some((x1, y1)) = step((x, y), i, size) {
                        if palette.at((x1, y1), size) == color &&
                            checked.at((x1, y1), size) == 0 {  // Not viewed before
                            checked.paint_at((x1, y1), n_patch, size);
                            queue.push_back((x1, y1));
                        }
                    }
                }
            }
        }
    }

    n_patch
}

/// Returns a vector of `(cells, neighbors)` where `cells` represents the content of
/// the patch with the same color and `neighbors` represents the neighboring colors of
/// that patch.
pub fn view(palette: &Palette<Color>, size: &PaletteSize)
    -> Vec<(Vec<(usize, usize)>, HashSet<Color>)> {
    let mut checked = Palette::new_const(0u32, size);
    let mut n_patch = 0;
    let mut result: Vec<(Vec<(usize, usize)>, HashSet<Color>)> = Vec::new();

    for i in 0..size.h {
        for j in 0..size.w {
            if checked.at((i, j), size) != 0 || palette.at((i, j), size) == BLANK { continue }
            // else
            n_patch += 1;

            let mut queue: VecDeque<(usize, usize)> = VecDeque::new();
            let mut cells: Vec<(usize, usize)> = Vec::new();
            let mut neighbors: HashSet<Color> = HashSet::new();  // Neighboring colors

            checked.paint_at((i, j), n_patch, size);
            queue.push_back((i, j));

            let color = palette.at((i, j), size);
            while !queue.is_empty() {
                let (x, y) = queue.pop_front().unwrap();
                cells.push((x, y));

                for i in 0..4 {
                    if let Some((x1, y1)) = step((x, y), i, size) {
                        let color2 = palette.at((x1, y1), size);
                        if color2 != BLANK {
                            if color2 != color {
                                neighbors.insert(color2);
                            } else if checked.at((x1, y1), size) == 0 {  // Not viewed before
                                checked.paint_at((x1, y1), n_patch, size);
                                queue.push_back((x1, y1));
                            }
                        }
                    }
                }
            }
            result.push((cells, neighbors));
        }
    }

    result
}
